# -*- coding:utf8 -*-
from __future__ import  print_function, unicode_literals
import ahocorasick
import pinyin

from logger_utility.loggconfig import getLogger


class AutomatonWrap(object):

    logger = getLogger('WordAutomation')

    def __init__(self, words):
        self.model = ahocorasick.Automaton()
        self.map = {}
        for word in words:
            self._add_word(word, word)
        self.model.make_automaton()

    def _add_word(self, k, v):
        """
        添加词语
        :param k: 
        :param v: 
        :return: 
        """
        self.model.add_word(k, v)
        self.map[v] = k

    def iter(self, content):
        """
        进行迭代并取最长的那个
        :param content: 
        :return: 
        """
        ret = []
        last_begin, last_end = None, None

        for i in self.model.iter(content):
            end_index, v = i
            begin_index = end_index - len(self.map[v]) + 1
            # self.logger.debug("Find Key Word %s, %s, %s", begin_index, v, end_index)

            if last_end:
                if last_begin >= begin_index and end_index >= last_end:  # [cb [lb le] ce]
                    ret[-1] = v
                    last_begin, last_end = begin_index, end_index
                    continue
                elif begin_index >= last_begin and end_index <= last_end: # [lb [cb ce] le]
                    continue

            ret.append(v)
            last_begin, last_end = begin_index, end_index

        return ret


class PinYinAutomatonWrap(object):
    """
    基于拼音的全文匹配
    """

    logger = getLogger('PinYinAutoMaton')

    def __init__(self, words):
        self.model = ahocorasick.Automaton()
        self.map = {}
        self.pinyin_map = {}

        added = set()
        for word in words:
            pinyin_ = pinyin.get(word, format='strip', delimiter='')
            if pinyin_ in added:
                continue
            self.pinyin_map[pinyin_] = word
            added.add(pinyin_)
            self._add_word(pinyin_, pinyin_)
        self.model.make_automaton()

    def _add_word(self, k, v):
        """
        添加词语
        :param k: 
        :param v: 
        :return: 
        """
        self.model.add_word(k, v)

    def _pinyin_2_word(self, pinyin_index, content, begin_index, end_index, pinyin_):
        """
        将拼音还原成汉子
        """
        try:
            begin_w = pinyin_index[begin_index]
            end_w = pinyin_index.get(end_index)
            if end_w:
                return content[begin_w: end_w]
            return None

        except Exception as e:
            self.logger.warning("Pinyin To Word Wrong Pinyin: %s E: %s", pinyin_, e)

    def iter(self, content):
        """
        进行迭代并取最长的那个
        :param content: 
        :return: 
        """
        pinyin_content = ''
        pinyin_index = {}
        for i, char in enumerate(content):
            pinyin_ = pinyin.get(char, format='strip')
            pinyin_index[len(pinyin_content)] = i
            pinyin_content += pinyin_  # + ''
        pinyin_index[len(pinyin_content)] = len(content)

        # self.logger.debug(sorted(pinyin_index.items()))

        ret = []
        last_begin, last_end = None, None
        for i in self.model.iter(pinyin_content):
            end_i, v = i
            end_ni = end_i + 1  # end，　不包含原字符
            begin_i = end_ni - len(v)
            word = self._pinyin_2_word(pinyin_index, content, begin_i, end_ni, v)
            self.logger.debug("Pinyin Senstive Word:%s, PinYin:%s, Origin:%s",
                              word,
                              v,
                              self.pinyin_map.get(v)
                              )

            if word:
                if last_end:
                    if last_begin >= begin_i and end_i >= last_end:  # [cb [lb le] ce]
                        ret[-1] = word
                        last_begin, last_end = begin_i, end_i
                        continue
                    elif begin_i >= last_begin and end_i <= last_end:  # [lb [cb ce] le]
                        continue

                ret.append(word)
                last_begin, last_end = begin_i, end_i

        return ret


if __name__ == '__main__':
    v = PinYinAutomatonWrap([u'一个人', u'er个人'])
    for i in v.iter(u'一个人二 个人'):
        print(i)
